from functools import partial

import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp

from utils.auto_docx import convert_to_docx
from utils.desicion_parameters import catch_mass as mass
from utils.desicion_parameters import elastic_modulus as em
from utils.desicion_parameters import force_external as force
from utils.desicion_parameters import force_tension_start as T0
from utils.desicion_parameters import length_loaded as length
from utils.desicion_parameters import section_area_wire as fi


def diff_equation_dim_1(time, w0: list):
    """Функция задания дифф. ур. движения

    ПАРАМЕТРЫ
    ---------
    - time : список моментов времени (см. подробнее в описании solve_ivp)
    - w : начальные условия (см. подробнее в описании solve_ivp)

    РЕЗУЛЬТАТЫ
    ----------
    - dw1 : первая производная координаты по времени
    в текущий момент времени, [m/s]
    - dw2 : вторая производная координаты по времени
    в текущий момент времени, [m/s^2]


    ДОПОЛНИТЕЛЬНО
    -------------
    - можно поиграться с переменной p_external
    для изменения вынуждающей силы
    """

    p_external = force * np.sin(2 * time**2) / mass
    dw1 = w0[1]
    dw2 = (
        p_external
        - (2 * T0 * w0[0]) / (length * mass)
        - (em * fi * (w0[0] / length) ** 3) / mass
    )
    return [dw1, dw2]


def solve_equation(time: float, time_div: int, w0: float, w1: float):
    """Решаeт дифф. ур. движения груза. Расчет по времени от 0
    до time.

    ПАРАМЕТРЫ
    ---------
    - time : конечный момент времени [s]
    - time_div : количество элементов, на которые будет разбит
    временной диапазон, б/р
    - w0 : начальная координата по Oz [m]
    - w1 : начальная скорость по Oz [m/s]

    РЕЗУЛЬТАТЫ
    ----------
    см. подробнее описание solve_ivp


    ДОПОЛНИТЕЛЬНО
    -------------
    - ничего не трогать!!!
    """

    return solve_ivp(
        fun=diff_equation_dim_1,
        t_span=[0, time],
        y0=[w0, w1],
        t_eval=[
            round(x, 1 + len(str(time_div))) for x in np.linspace(0, time, time_div)
        ],
        atol=10 ** (-10),
        rtol=10 ** (-10),
        dense_output=True,
    )


def coord_from_time(time_coord, coord, show: bool):
    """Функция выводит график координаты от времени

    ПАРАМЕТРЫ
    ---------
    - time_coord : список моментов времени из решения уравнения
    - coord : список положения грузика из решения уравнения
    - show : определяет, покажет ли график во время выполнения программы

    РЕЗУЛЬТАТЫ
    ----------
    - ./utils/w_m.png : вывод результатов в файл
    """

    fig, ax = plt.subplots()
    ax.plot(time_coord, coord)
    plt.xlabel("t, s")
    plt.ylabel("W, m")
    plt.legend(["W, m"], shadow=True)
    plt.title("W(t)")
    fig.savefig(
        "dz/non_linear_1d_vibration/utils/w_m.png",
        facecolor="w",
        bbox_inches="tight",
        pad_inches=0.1,
        transparent=True,
    )
    if show:
        plt.show()


def animate_phase(
    i, line=None, coord=None, velocity=None, text=None, time=None, text2=None
):
    """Задает прорисовку кадра фазового портрета системы
    ПАРАМЕТРЫ
    ---------
    - i : момент из frame (см. matplotlib.FuncAnimation)
    - line : линия - следующий шаг фазового портрета
    - coord : список координат грузика из решения уравнения
    - velocity : список скоростей грузика из решения уравнения
    - text : вывод параметров этого момента колебательной системы
    - time : список моментов времени грузика из решения уравнения
    - text2 : подчеркивание этого момента колебательной системы

    РЕЗУЛЬТАТЫ
    ----------
    - см. matplotlib.FuncAnimation

    """
    line.axes.axis(
        [
            1.30 * min(coord),
            1.30 * max(coord),
            1.30 * min(velocity),
            1.30 * max(velocity),
        ]
    )
    line.set_data(coord[:i], velocity[:i])
    text.set_position((1.25 * min(coord), 1.25 * min(velocity)))
    text.set_text(
        f"""Time = {round(time[i], 5)
                    } s\nX = {round(coord[i], 5)
                    } m\nX' = {round(velocity[i], 5)} m/s"""
    )
    text2.set_position((coord[i], velocity[i]))
    text2.set_text(f"{round(time[i],5)} s")
    return (line,)


def phase_portrait(time, coord, velocity, frm: int, show: bool):
    """Функция анимирует фазовый портрет скорости от координаты

    ПАРАМЕТРЫ
    ---------
    - time : список моментов времени из решения уравнения
    - coord : список положения грузика из решения уравнения
    - velocity : список скоростей грузика из решения уравнения
    - frm : определяет количество записанных кадров в gif
    должно быть не больше длины списка time
    - show : определяет, покажет ли график во время выполнения программы

    РЕЗУЛЬТАТЫ
    ----------
    - ./utils/phase_portrait.gif : вывод результатов в файл
    """

    fig, ax = plt.subplots()
    (line,) = ax.plot(coord, velocity)
    text = ax.text(100, 100, "", color="k")
    text2 = ax.text(1, 1, "", color="b")
    plt.xlabel("w, m")
    plt.ylabel("w', m/s")
    plt.legend(["w'(w)"], shadow=True)
    plt.title("Фазовый портрет")

    ani = animation.FuncAnimation(
        fig,
        partial(
            animate_phase,
            line=line,
            coord=coord,
            velocity=velocity,
            text=text,
            time=time,
            text2=text2,
        ),
        frames=frm,
        interval=100,
        repeat=False,
    )
    ani.save(
        "dz/non_linear_1d_vibration/utils/phase_portrait.gif",
        writer="imagemagick",
        fps=30,
    )
    if show:
        plt.show()


def animate_syst(i, ax=None, coord=None, time=None):
    """Задает прорисовку кадра колебания системы
    ПАРАМЕТРЫ
    ---------
    - i : момент из frame (см. matplotlib.FuncAnimation)
    - ax : окно графического вывода
    - coord : список координат грузика из решения уравнения
    - time : список моментов времени грузика из решения уравнения

    РЕЗУЛЬТАТЫ
    ----------
    - см. matplotlib.FuncAnimation

    """
    ax.clear()
    plt.xlabel("x, m")
    plt.ylabel("y, m")
    plt.title("Анимация колебания")
    text = ax.text(100, 100, "", color="k")
    text.set_text(
        f"""Time = {round(time[i],5)
                    } s\nX = {round(coord[i],5)} m"""
    )
    text.set_position((-length * 1.15, min(coord) * 1.45))
    point = [0, coord[i]]
    line1_x = [-length, 0]
    line2_x = [length, 0]
    ax.plot(line1_x, point, color="b", label="original", marker="o")
    ax.plot(line2_x, point, color="b", label="original", marker="o")
    ax.plot(point[0], point[1], color="green", label="original", marker="o")
    ax.set_xlim([-length * 1.2, length * 1.2])
    ax.set_ylim([1.5 * min(coord), 1.5 * max(coord)])
    return


def catcher_anim(time: list, coord: list, frm: int, show: bool):
    """Функция анимирует колебания грузика.

    ПАРАМЕТРЫ
    ---------
    - time : список моментов времени из решения уравнения
    - coord : список положения грузика из решения уравнения
    - frm : определяет количество записанных кадров в gif
    должно быть не больше длины списка time
    - show : определяет, покажет ли график во время выполнения программы

    РЕЗУЛЬТАТЫ
    ----------
    - ./utils/animation.gif : вывод результатов в файл
    """

    fig, ax = plt.subplots()
    ani = animation.FuncAnimation(
        fig,
        partial(
            animate_syst,
            ax=ax,
            coord=coord,
            time=time,
        ),
        frames=frm,
        interval=100,
        repeat=False,
    )
    ani.save(
        "dz/non_linear_1d_vibration/utils/animation.gif", writer="imagemagick", fps=30
    )
    if show:
        plt.show()


sol = solve_equation(time=2, time_div=10000, w0=0.012, w1=0)

coord_from_time(time_coord=sol.t, coord=sol.y[0], show=False)
phase_portrait(time=sol.t, coord=sol.y[0], velocity=sol.y[1], frm=700, show=False)
catcher_anim(time=sol.t, coord=sol.y[0], frm=300, show=False)
convert_to_docx()
